import torch
from .chamfer_distance import chamfer_distance


def get_line_instances_from_mask(mask, scale_x, scale_y):
    # mask: H, W
    # instance: {1: (N1, 2), 2: (N2, 2), ..., k1: (N_k1, 2)}
    indices = torch.unique(mask)
    instances = {}
    for idx in indices:
        idx = idx.item()
        if idx == 0:
            continue
        pc_x, pc_y = torch.where(mask == idx)
        pc_x = pc_x.float() * scale_x
        pc_y = pc_y.float() * scale_y
        coords = torch.stack([pc_x, pc_y], -1)
        instances[idx] = coords
    return instances

def get_line_seg_from_mask(mask):
    # mask: H, W
    # instance: {1: (N1, 2), 2: (N2, 2), ..., k1: (N_k1, 2)}
    indices = torch.unique(mask)
    instances = {}
    for idx in indices:
        idx = idx.item()
        if idx == 0:
            continue
        pc_x, pc_y = torch.where(mask == idx)
        coords = torch.stack([pc_x, pc_y], -1)
        instances[idx] = coords
    return instances

def line_matching_by_CD(inst_pred_lines, 
                        inst_pred_confidence, 
                        inst_label_lines, 
                        seg_pred_lines, 
                        seg_label_lines, 
                        threshold, 
                        bidirectional=False, 
                        threshold_iou=0.1):
    # inst_pred_line: a list of points {1: (M1, 2), 2: (M2, 2), ..., k2: (M_k2, 2)}
    # inst_pred_confidence: a list of confidence [c1, c2, ..., ck2]
    # inst_label_line: a list of points {1: (N1, 2), 2: (N2, 2), ..., k1: (N_k1, 2)}
    # return: a list of {'pred': (M, 2), 'label': (N, 2), 'confidence': scalar}
    pred_num = len(inst_pred_lines)
    label_num = len(inst_label_lines)
    CD = torch.zeros((pred_num, label_num)).cuda()
    iou = torch.zeros((pred_num, label_num)).cuda()

    inst_pred_lines_keys = [*inst_pred_lines]
    inst_label_lines_keys = [*inst_label_lines]
    for i, key_pred in enumerate(inst_pred_lines_keys):
        for j, key_label in enumerate(inst_label_lines_keys):
            CD[i, j] = chamfer_distance(inst_pred_lines[key_pred][None], inst_label_lines[key_label][None], bidirectional=bidirectional, threshold=threshold)

            seg_pred_line = seg_pred_lines[key_pred][None][0]
            seg_label_line = seg_label_lines[key_label][None][0]
            seg_pred = torch.zeros(600,800)
            seg_label = torch.zeros(600,800)
            seg_pred[seg_pred_line[:, 0], seg_pred_line[:, 1]] = 1
            seg_label[seg_label_line[:, 0], seg_label_line[:, 1]] = 1
            pred_map = seg_pred.bool()
            gt_map = seg_label.bool()
            intersect = (pred_map & gt_map).sum().float()
            union = (pred_map | gt_map).sum().float()
            iou[i, j] = intersect / (union + 1e-7)

    pred_taken = torch.zeros(pred_num, dtype=torch.bool).cuda()
    label_taken = torch.zeros(label_num, dtype=torch.bool).cuda()
    matched_list = []
    if pred_num > 0 and label_num > 0:
        while True:
            idx = torch.argmin(CD)
            i, j = (idx // CD.shape[1]).item(), (idx % CD.shape[1]).item()
            if CD[i, j] >= threshold:
                break

            if iou[i, j] >= threshold_iou:
                matched_list.append({
                    'pred': inst_pred_lines[inst_pred_lines_keys[i]],
                    'confidence': inst_pred_confidence[inst_pred_lines_keys[i]],
                    'label': inst_label_lines[inst_label_lines_keys[j]],
                    'CD': CD[i, j].item(),
                })
                pred_taken[i] = True
                label_taken[j] = True
            CD[i, :] = threshold
            CD[:, j] = threshold

    for i in range(pred_num):
        if not pred_taken[i]:
            matched_list.append({
                'pred': inst_pred_lines[inst_pred_lines_keys[i]],
                'confidence': inst_pred_confidence[inst_pred_lines_keys[i]],
                'label': None,
                'CD': threshold,
            })

    for j in range(label_num):
        if not label_taken[j]:
            matched_list.append({
                'pred': None,
                'confidence': 0,
                'label': inst_label_lines[inst_label_lines_keys[j]],
                'CD': threshold,
            })

    return matched_list


def instance_mask_AP(AP_matrix, 
                     AP_count_matrix, 
                     inst_pred_mask, 
                     inst_label_mask, 
                     scale_x, 
                     scale_y, 
                     confidence, 
                     thresholds, 
                     sampled_recalls, 
                     bidirectional=False, 
                     threshold_iou=0.1):
    # inst_pred: N, C, H, W
    # inst_label: N, C, H, W
    # confidence: N, max_instance_num
    N, C, H, W = inst_label_mask.shape
    for n in range(N):
        for c in range(C):
            inst_pred_lines = get_line_instances_from_mask(inst_pred_mask[n, c], scale_x, scale_y)
            inst_label_lines = get_line_instances_from_mask(inst_label_mask[n, c], scale_x, scale_y)

            seg_pred_lines = get_line_seg_from_mask(inst_pred_mask[n, c])
            seg_label_lines = get_line_seg_from_mask(inst_label_mask[n, c]) # 区别是像素坐标，上述为bev空间坐标

            if len(inst_pred_lines) == 0 and len(inst_label_lines) == 0:
                continue
            AP_matrix[c] += single_instance_line_AP(inst_pred_lines, confidence[n], inst_label_lines, 
                                                    thresholds, seg_pred_lines, seg_label_lines, 
                                                    sampled_recalls=sampled_recalls, 
                                                    bidirectional=bidirectional, 
                                                    threshold_iou=threshold_iou)
            AP_count_matrix[c] += 1


def single_instance_line_AP(inst_pred_lines, 
                            inst_pred_confidence, 
                            inst_label_lines, 
                            thresholds, 
                            seg_pred_lines, 
                            seg_label_lines, 
                            sampled_recalls, 
                            bidirectional=False, 
                            threshold_iou=0.1):
    # inst_pred_line: a list of points {1: (M1, 2), 2: (M2, 2), ..., k2: (M_k2, 2)}
    # inst_pred_confidence: a list of confidence [c1, c2, ..., ck2]
    # inst_label_line: a list of points {1: (N1, 2), 2: (N2, 2), ..., k1: (N_k1, 2)}
    # thresholds: threshold of chamfer distance to identify TP
    num_thres = len(thresholds)
    AP_thres = torch.zeros(num_thres).cuda()
    for t in range(num_thres):
        matching_list = line_matching_by_CD(inst_pred_lines, inst_pred_confidence, 
                                            inst_label_lines, seg_pred_lines, seg_label_lines, 
                                            thresholds[t], bidirectional=bidirectional, 
                                            threshold_iou=threshold_iou)
        precision, recall = get_precision_recall_curve_by_confidence(matching_list, len(inst_label_lines), thresholds[t])
        precision, recall = smooth_PR_curve(precision, recall)
        AP = calc_AP_from_precision_recall(precision, recall, sampled_recalls)
        AP_thres[t] = AP
    return AP_thres


def get_precision_recall_curve_by_confidence(matching_list, num_gt, threshold):
    matching_list = sorted(matching_list, key=lambda x: x['confidence'])

    TP = [0]
    FP = [0]
    for match_item in matching_list:
        pred = match_item['pred']
        label = match_item['label']
        dist = match_item['CD']

        if pred is None:
            continue

        if label is None:
            TP.append(TP[-1])
            FP.append(FP[-1] + 1)
            continue

        if dist < threshold:
            TP.append(TP[-1] + 1)
            FP.append(FP[-1])
        else:
            TP.append(TP[-1])
            FP.append(FP[-1] + 1)

    TP = torch.tensor(TP[1:])
    FP = torch.tensor(FP[1:])

    precision = TP / (TP + FP)
    recall = TP / num_gt
    return precision, recall


def smooth_PR_curve(precision, recall):
    idx = torch.argsort(recall)
    recall = recall[idx]
    precision = precision[idx]
    length = len(precision)
    for i in range(length-1, 0, -1):
        precision[:i][precision[:i] < precision[i]] = precision[i]
    return precision, recall


def calc_AP_from_precision_recall(precision, recall, sampled_recalls):
    acc_precision = 0.
    total = len(sampled_recalls)
    for r in sampled_recalls:
        idx = torch.where(recall >= r)[0]
        if len(idx) == 0:
            continue
        idx = idx[0]
        acc_precision += precision[idx]
    return acc_precision / total



def instance_mask_turn_cal(inst_label_mask, scale_x, scale_y, width_thr=6):
    # inst_pred: N, C, H, W
    # inst_label: N, C, H, W
    # confidence: N, max_instance_num
    # instance: {1: (N1, 2), 2: (N2, 2), ..., k1: (N_k1, 2)}
    N, C, H, W = inst_label_mask.shape
    for n in range(N):
        # for c in range(C):
        c = 2
        inst_label_lines = get_line_instances_from_mask(inst_label_mask[n, c], scale_x, scale_y)
        # print("len(inst_label_lines): ", len(inst_label_lines))
        if len(inst_label_lines) != 0:
            keys = list(inst_label_lines.keys())
            for i in range(len(inst_label_lines)):
                linex = inst_label_lines[keys[i]][:,0]
                # print(linex)
                width = torch.max(linex) - torch.min(linex)
                if width > width_thr:
                    return True
            return False
        else:
            return False

            